from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.config import get_store
from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore
from langchain_ollama import ChatOllama

llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

# case 1: tool with short memory
# store = InMemoryStore() 

# store.put(  
#     ("users",),  
#     "user_123",  
#     {
#         "name": "John Smith",
#         "language": "English",
#     } 
# )

# @tool
# def get_user_info(config: RunnableConfig) -> str:
#     """Look up user info."""
#     # Same as that provided to `create_react_agent`
#     store = get_store() 
#     user_id = config["configurable"].get("user_id")
#     user_info = store.get(("users",), user_id) 
#     return str(user_info.value) if user_info else "Unknown user"

# agent = create_react_agent(
#     model=llm,
#     tools=[get_user_info],
#     store=store 
# )

# # Run the agent
# result = agent.invoke(
#     {"messages": [{"role": "user", "content": "look up user information"}]},
#     config={"configurable": {"user_id": "user_123"}}
# )
# print(result)

# case 2: tool with long memory
from typing_extensions import TypedDict
from langchain_core.tools import tool
from langgraph.config import get_store
from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore

store = InMemoryStore() 

class UserInfo(TypedDict): 
    name: str

@tool
def save_user_info(user_info: UserInfo, config: RunnableConfig) -> str: 
    """Save user info."""
    # Same as that provided to `create_react_agent`
    store = get_store() 
    user_id = config["configurable"].get("user_id")
    store.put(("users",), user_id, user_info) 
    return "Successfully saved user info."

agent = create_react_agent(
    model=llm,
    tools=[save_user_info],
    store=store
)

# Run the agent
agent.invoke(
    {"messages": [{"role": "user", "content": "Save my name is John Smith"}]},
    config={"configurable": {"user_id": "user_123"}} 
)

# You can access the store directly to get the value
print(store.get(("users",), "user_123").value)
